import os
import json
import torch
import pandas as pd
from param_decoder import train_param_predictor, predict_params

def load_gplearn_results(results_dir="gplearn_results", prefix="Nguyen-11"):
    """加载GPlearn实验结果"""
    # 获取所有匹配的文件
    all_data = []
    
    # 遍历目录下的所有文件
    for file in os.listdir(results_dir):
        if file.startswith(prefix) and file.endswith('.csv'):
            file_path = os.path.join(results_dir, file)
            print(f"加载文件: {file}")
            
            try:
                df = pd.read_csv(file_path)
                
                # 提取MSE和参数数据
                for _, row in df.iterrows():
                    mse = row['mse']
                    params = row['params']
                    if isinstance(params, str):
                        params = eval(params)  # 将字符串转换为字典
                        params = {
                            "p_crossover": params['p_crossover'],
                            "p_subtree_mutation": params['p_subtree_mutation'],
                            "p_hoist_mutation": params['p_hoist_mutation'],
                            "p_point_mutation": params['p_point_mutation']
                        }
                    all_data.append((mse, params))
            except Exception as e:
                print(f"处理文件 {file} 时出错: {str(e)}")
                continue
    
    if not all_data:
        raise FileNotFoundError(f"在 {results_dir} 目录下没有找到前缀为 {prefix} 的结果文件")
    
    print(f"总共加载了 {len(all_data)} 条数据")
    return all_data

def main():
    # 设置设备
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    # 加载GPlearn实验结果
    print("加载GPlearn实验结果...")
    train_data = load_gplearn_results()
    print(f"加载了 {len(train_data)} 条数据")
    
    # 训练参数预测器
    print("\n开始训练参数预测器...")
    predictor = train_param_predictor(
        train_data,
        n_epochs=1000,
        batch_size=32,
        learning_rate=1e-4,
        device=device
    )
    
    # 保存模型
    save_dir = "param_predictor"
    os.makedirs(save_dir, exist_ok=True)
    
    torch.save(predictor.model.state_dict(), os.path.join(save_dir, "model.pth"))
    print(f"\n模型已保存到: {save_dir}")
    
    # 测试模型
    print("\n测试模型预测...")
    test_mse = 0.1  # 示例MSE值
    predicted_params = predict_params(
        predictor,
        test_mse,
        n_samples=5
    )
    
    print(f"\n对于MSE={test_mse}的预测参数:")
    for i, params in enumerate(predicted_params):
        print(f"样本 {i+1}:")
        for key, value in params.items():
            print(f"  {key}: {value:.4f}")

if __name__ == "__main__":
    main() 